import os
import os.path as osp
import logging
import yaml
from utils.util import OrderedYaml
Loader, Dumper = OrderedYaml()


def parse(src_path, dst_path, is_eval=True, is_fast = False):
    is_train = False

    # with open(opt_path, mode='r') as f:
    #     opt = yaml.load(f, Loader=Loader)
    opt = {}
    dataset = {}
    opt['path'] = {}
    opt['network_G'] = {}

    opt['path']['root'] = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir, osp.pardir))
    pre_trained_path = osp.join(opt['path']['root'],'pre_trained')
    if is_fast:
        opt['name'] = 'FSRCNN'
        opt['network_G'] = {
            'which_model_G':'classSR_3class_fsrcnn_net',
            'in_nc': 3,
            'out_nc': 3
        }
        opt['path']['pretrain_model_G'] =  osp.join(pre_trained_path,'ClassSR_FSRCNN.pth')
    else:
        opt['name'] = 'RCAN'
        opt['network_G'] = {
            'which_model_G':'classSR_3class_rcan',
            'in_nc': 3,
            'out_nc': 3
        }
        opt['path']['pretrain_model_G'] =  osp.join(pre_trained_path,'ClassSR_RCAN.pth')

    opt['is_train'] = is_train
    opt['scale'] = 4
    scale = opt['scale']
    opt['model'] = 'ClassSR'
    opt['patch_size'] = 32
    opt['step'] = 20
    opt['gpu_ids'] = []


    opt['network_G']['scale'] = scale

    # export CUDA_VISIBLE_DEVICES
    gpu_list = ','.join(str(x) for x in opt['gpu_ids'])
    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
    # print('export CUDA_VISIBLE_DEVICES=' + gpu_list)

    if is_eval:
        dataset['phase'] = 'test_eval'
        dataset['name'] = 'eval'
        dataset['scale'] = scale
        dataset['data_type'] = 'img'
        if osp.isdir(src_path):
            dataset['dataroot_GT'] = src_path
        else:
            dataset['dataroot_GT'] = osp.join(opt['path']['root'],'temp','GT')
        dataset['dataroot_LQ'] = osp.join(opt['path']['root'],'temp','LQ')
        dataset['mode'] = 'LQGT'

        opt['path']['results_root'] = osp.join(opt['path']['root'],'temp','result')
        opt['path']['log'] = osp.join(opt['path']['root'], 'log')


    else :
        dataset['phase'] = 'test_generate'
        dataset['name'] = 'generate'
        dataset['scale'] = scale
        dataset['data_type'] = 'img'
        if osp.isdir(src_path):
            dataset['dataroot_LQ'] = src_path
        else:
            dataset['dataroot_LQ'] = osp.join(opt['path']['root'],'temp','LQ')
        dataset['mode'] = 'LQ'

        opt['path']['results_root'] = dst_path
        opt['path']['log'] = osp.join(opt['path']['root'], 'log')



    return opt,dataset



def dict2str(opt, indent_l=1):
    '''dict to string for logger'''
    msg = ''
    for k, v in opt.items():
        if isinstance(v, dict):
            msg += ' ' * (indent_l * 2) + k + ':[\n'
            msg += dict2str(v, indent_l + 1)
            msg += ' ' * (indent_l * 2) + ']\n'
        else:
            msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
    return msg


class NoneDict(dict):
    def __missing__(self, key):
        return None


# convert to NoneDict, which return None for missing key.
def dict_to_nonedict(opt):
    if isinstance(opt, dict):
        new_opt = dict()
        for key, sub_opt in opt.items():
            new_opt[key] = dict_to_nonedict(sub_opt)
        return NoneDict(**new_opt)
    elif isinstance(opt, list):
        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
    else:
        return opt


def check_resume(opt, resume_iter):
    '''Check resume states and pretrain_model paths'''
    logger = logging.getLogger('base')
    if opt['path']['resume_state']:
        if opt['path'].get('pretrain_model_G', None) is not None or opt['path'].get(
                'pretrain_model_D', None) is not None:
            logger.warning('pretrain_model path will be ignored when resuming training.')

        opt['path']['pretrain_model_G'] = osp.join(opt['path']['models'],
                                                   '{}_G.pth'.format(resume_iter))
        logger.info('Set [pretrain_model_G] to ' + opt['path']['pretrain_model_G'])
        if 'gan' in opt['model']:
            opt['path']['pretrain_model_D'] = osp.join(opt['path']['models'],
                                                       '{}_D.pth'.format(resume_iter))
            logger.info('Set [pretrain_model_D] to ' + opt['path']['pretrain_model_D'])
